#!/usr/bin/env python3
"""
Milvus 入门 Demo
演示基本的向量数据库操作：
1. 连接到Milvus
2. 创建集合
3. 插入向量数据
4. 执行向量搜索
5. 查询和删除数据
"""

import random
import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
)

# 配置参数
COLLECTION_NAME = "hello_milvus"
DIM = 128  # 向量维度
NUM_ENTITIES = 1000  # 实体数量

def connect_to_milvus():
    """连接到Milvus服务器"""
    print("正在连接到Milvus...")
    try:
        # 连接到本地Milvus服务器（默认端口19530）
        connections.connect(
            alias="default",
            host="localhost",
            port="19530"
        )
        print("✅ 成功连接到Milvus!")
        return True
    except Exception as e:
        print(f"❌ 连接失败: {e}")
        print("请确保Milvus服务器正在运行")
        return False

def create_collection():
    """创建集合"""
    print(f"\n正在创建集合 '{COLLECTION_NAME}'...")
    
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="age", dtype=DataType.INT64),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIM)
    ]
    
    # 创建集合模式
    schema = CollectionSchema(
        fields=fields,
        description="Hello Milvus 演示集合"
    )
    
    # 创建集合
    collection = Collection(
        name=COLLECTION_NAME,
        schema=schema,
        using="default"
    )
    
    # 创建索引
    print("正在创建索引...")
    index_params = {
        "metric_type": "L2",  # 距离度量类型
        "index_type": "IVF_FLAT",  # 索引类型
        "params": {"nlist": 128}  # 聚类数量
    }
    collection.create_index(
        field_name="embedding",
        index_params=index_params
    )
    
    print("✅ 集合创建成功!")
    return collection

def insert_data(collection):
    """插入数据"""
    print(f"\n正在插入 {NUM_ENTITIES} 条数据...")
    
    # 生成随机数据
    ages = [random.randint(18, 80) for _ in range(NUM_ENTITIES)]
    embeddings = [[random.random() for _ in range(DIM)] for _ in range(NUM_ENTITIES)]
    
    # 准备插入数据
    entities = [
        ages,
        embeddings
    ]
    
    # 插入数据
    collection.insert(entities)
    collection.flush()  # 确保数据持久化
    
    print(f"✅ 成功插入 {NUM_ENTITIES} 条数据!")

def search_vectors(collection):
    """执行向量搜索"""
    print("\n正在执行向量搜索...")
    
    # 加载集合到内存
    collection.load()
    
    # 生成搜索向量
    search_vectors = [[random.random() for _ in range(DIM)]]
    
    # 搜索参数
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10}
    }
    
    # 执行搜索
    results = collection.search(
        data=search_vectors,
        anns_field="embedding",
        param=search_params,
        limit=5,  # 返回前5个最相似的结果
        output_fields=["age"]  # 同时返回age字段
    )
    
    print("🔍 搜索结果:")
    for i, hits in enumerate(results):
        print(f"查询向量 {i + 1}:")
        for j, hit in enumerate(hits):
            print(f"  第 {j + 1} 名: ID={hit.id}, 距离={hit.distance:.4f}, 年龄={hit.entity.get('age')}")

def query_data(collection):
    """查询数据"""
    print("\n正在查询数据...")
    
    # 查询所有数据
    results = collection.query(
        expr="age > 30",
        output_fields=["id", "age"],
        limit=10
    )
    
    print("📊 查询结果 (年龄 > 30 的前10条):")
    for result in results:
        print(f"  ID: {result['id']}, 年龄: {result['age']}")

def get_collection_stats(collection):
    """获取集合统计信息"""
    print("\n📈 集合统计信息:")
    print(f"  集合名称: {collection.name}")
    print(f"  实体数量: {collection.num_entities}")
    print(f"  字段数量: {len(collection.schema.fields)}")
    
    # 显示字段信息
    print("  字段详情:")
    for field in collection.schema.fields:
        print(f"    - {field.name}: {field.dtype}")

def cleanup():
    """清理资源"""
    print("\n🧹 清理资源...")
    
    # 删除集合
    if utility.has_collection(COLLECTION_NAME):
        utility.drop_collection(COLLECTION_NAME)
        print(f"✅ 已删除集合 '{COLLECTION_NAME}'")
    
    # 断开连接
    connections.disconnect("default")
    print("✅ 已断开Milvus连接")

def main():
    """主函数"""
    print("🚀 Milvus 入门 Demo")
    print("=" * 50)
    
    # 连接到Milvus
    if not connect_to_milvus():
        return
    
    try:
        # 检查集合是否已存在
        if utility.has_collection(COLLECTION_NAME):
            print(f"集合 '{COLLECTION_NAME}' 已存在，正在删除...")
            utility.drop_collection(COLLECTION_NAME)
        
        # 创建集合
        collection = create_collection()
        
        # 插入数据
        insert_data(collection)
        
        # 获取统计信息
        get_collection_stats(collection)
        
        # 执行搜索
        search_vectors(collection)
        
        # 查询数据
        query_data(collection)
        
        print("\n🎉 Demo 执行完成!")
        
    except Exception as e:
        print(f"❌ 执行过程中出现错误: {e}")
    
    finally:
        # 清理资源
        cleanup()

if __name__ == "__main__":
    main()
